/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Contains classes describing angular distribution functions.

// $Id$

#ifndef __angular__
#define __angular__

#include "mcrx-types.h"
#include "ray.h"
#include "constants.h"
#include "boost/shared_ptr.hpp"
#include "blitz/array.h"
#include "mono_poly_abstract.h"

namespace mcrx {
  template <typename> class angular_distribution;
  template<typename> class isotropic_angular_distribution;
  template <typename T> class isotropic_angular_distribution<blitz::Array<T, 1> >; 
  template<typename> class collimated_angular_distribution;
  template <typename T> class collimated_angular_distribution<blitz::Array<T, 1> >; 
  template<typename> class cone_angular_distribution;
  template <typename T> class cone_angular_distribution<blitz::Array<T,1> >;
  template<typename> class ray_distribution_pair;
}


/** Abstract base class for objects that can tell the transfer class
    what the probability is of emission/scattering in certain angles.
    This is necessary for using the next-event estimator.
*/
template <typename T_lambda> 
class mcrx::angular_distribution {
public:
  angular_distribution () {};
  virtual ~angular_distribution () {};
  /** This returns the distribution function. When T_lambda is a blitz
      array, this is less efficient because it requires an
      allocation. However, we can't return an expression because it's
      a virtual function. */
  virtual T_lambda distribution_function(const vec3d& new_direction,
					 const vec3d& old_direction) const= 0;
  /** This function assigns the distribution function in place. This
      is more efficient for arrays since we don't have to allocate the
      return value. */
  virtual void distribution_function(const vec3d& new_direction,
				     const vec3d& old_direction,
				     T_lambda& distr) const= 0;
  virtual boost::shared_ptr<angular_distribution<T_lambda> > clone() const=0;
};


/** Angular_distribution class for isotropic distributions.  Simply
    returns 1/4pi as the distribution function. Used by most of the
    general emission classes. */
template <typename T_lambda> 
class mcrx::isotropic_angular_distribution: 
  public angular_distribution<T_lambda>
{
public:
  isotropic_angular_distribution (const T_lambda& l): 
    angular_distribution<T_lambda>() {};
  T_lambda distribution_function (const vec3d&, const vec3d&) const {
    return 1/(4*constants::pi); };
  void distribution_function(const vec3d&, const vec3d&,
			     T_lambda& distr) const {
    distr = 1/(4*constants::pi); };
  boost::shared_ptr<angular_distribution<T_lambda> > clone() const {
    return boost::shared_ptr<angular_distribution<T_lambda> >
      (new isotropic_angular_distribution<T_lambda>(*this));};
};

/** Specialization of the isotropic_angular_distribution class for
    polychromatic situations. */
template <typename T> 
class mcrx::isotropic_angular_distribution<blitz::Array<T,1> >: 
  public angular_distribution<blitz::Array<T,1> >
{
protected:
  typedef blitz::Array<T,1> T_lambda;

  /// The length of the wavelength vector.
  const int n;

public:
  isotropic_angular_distribution (const T_lambda& l):
    angular_distribution<T_lambda> (), n(l.size()) {};

  /** Returns the distribution function for an isotropic distribution,
      which is just 1/4pi.  To save space, we do not keep an array of
      the constant number, we create the array on-the-fly each time.
      Both alternatives are bad, but it is unfortunately not possible
      to return a constant since we are overloading a virtual
      function.  */
  T_lambda distribution_function (const vec3d&, const vec3d&) const {
    T_lambda d(n);
    isotropic_angular_distribution::distribution_function(vec3d(), vec3d(), d);
    return d;};
  void distribution_function (const vec3d&, const vec3d&,
			      T_lambda& distr) const {
    distr =  1./(4*constants::pi); };
  boost::shared_ptr<angular_distribution<blitz::Array<T,1> > > clone() const {
    return boost::shared_ptr<angular_distribution<blitz::Array<T,1> > >
      (new isotropic_angular_distribution<blitz::Array<T,1> >(*this));};
};


/** Angular_distribution class for collimated rays.  Simply returns 1
    for the set direction, otherwise 0. (This is not really valid, as
    it should be a delta function, but in any case the distribution is
    nonzero only on a set of measure zero, so it can't really be
    sampled with MC techniques.) */
template <typename T_lambda> 
class mcrx::collimated_angular_distribution:
  public angular_distribution<T_lambda>
{
  const vec3d direction_;

public:
  collimated_angular_distribution (const vec3d& dir, const T_lambda& l):
    angular_distribution<T_lambda> (), direction_ (dir) {};

  const vec3d& direction () const {return direction_;};
  /// The distribution function.
  T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    return (dot (dir, direction_) - 1) >-1e-9?  1:0;};
  T_lambda distribution_function (const vec3d& dir, const vec3d&,
				  T_lambda& distr) const {
    distr = (dot (dir, direction_) - 1) >-1e-9?  1:0; };
  boost::shared_ptr<angular_distribution<T_lambda> > clone() const {
    return boost::shared_ptr<angular_distribution<T_lambda> >
      (new collimated_angular_distribution<T_lambda>(*this));};
};

/** Specialization of the collimated_angular_distribution class for
    polychromatic situations. */
template <typename T> 
class mcrx::collimated_angular_distribution<blitz::Array<T,1> >: 
  public angular_distribution<blitz::Array<T,1> >
{
  typedef blitz::Array<T,1> T_lambda;

  /// The length of the wavelength vector.
  const int n;
  const vec3d direction_;

public:
  collimated_angular_distribution (const vec3d& dir, const T_lambda& l):
    angular_distribution<T_lambda> (), direction_ (dir), n (l.size())  {};

  const vec3d& direction () const {return direction_;};
  /// The distribution function.
  T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    T_lambda d(n); 
    d= collimated_angular_distribution::distribution_function(dir, vec3d(), d);
    return d;};
  void distribution_function (const vec3d& dir, const vec3d&,
			      T_lambda& distr) const {
    distr = (dot (dir, direction_) - 1) >-1e-9?  1:0; };

  boost::shared_ptr<angular_distribution<blitz::Array<T,1> > > clone() const {
    return boost::shared_ptr<angular_distribution<blitz::Array<T,1> > >
      (new collimated_angular_distribution<blitz::Array<T,1> >(*this));};
};


/** Angular_distribution class for an isotropic distribution within a
    certain angle around a specified direction. */
template <typename T_lambda> 
class mcrx::cone_angular_distribution: 
  public angular_distribution<T_lambda>
{
private:
  const vec3d direction_;
  T_float cos_theta_; ///< Cos of the opening angle.
  T_float value_; ///< The distribution value is 1./total solid angle subtended.

public:
  cone_angular_distribution (const vec3d& dir, T_float theta,
			     const T_lambda& l): 
    angular_distribution<T_lambda>(), 
    direction_(dir), 
    cos_theta_(cos(theta)), 
    value_(1./(4*constants::pi*sin(theta/2)*sin(theta/2))) {};
  const vec3d& direction () const {return direction_;};
  /** Returns the distribution function for a cone angular
      distribution, which is 1/omega, where omega is the total solid
      angle subtended by the cone, if we are within the cone and zero
      otherwise. */
  T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    const T_float costheta = dot(dir,direction_);
    return costheta > cos_theta_ ? value_ : 0;
  };
  T_lambda distribution_function (const vec3d& dir, const vec3d&,
				  T_lambda& distr) const {
    const T_float costheta = dot(dir,direction_);
    distr = (costheta > cos_theta_) ? value_ : 0;
  };

  boost::shared_ptr<angular_distribution<T_lambda> > clone() const {
    return boost::shared_ptr<angular_distribution<T_lambda> >
      (new cone_angular_distribution<T_lambda>(*this));};
};


/** Specialization of the cone_angular_distribution class for
    polychromatic situations. */
template <typename T> 
class mcrx::cone_angular_distribution<blitz::Array<T,1> >: 
  public angular_distribution<blitz::Array<T,1> >
{
protected:
  typedef blitz::Array<T,1> T_lambda;
  const vec3d direction_;
  T_float cos_theta_; ///< Cos of the opening angle.
  T_float value_; ///< The distribution value is 1./total solid angle subtended.

  /// The length of the wavelength vector.
  const int n;

public:
  cone_angular_distribution (const vec3d& dir, T_float theta,
			     const T_lambda& l): 
    angular_distribution<T_lambda> (), 
    direction_(dir), 
    cos_theta_(cos(theta)), 
    value_(1./(4*constants::pi*sin(theta/2)*sin(theta/2))),
    n(l.size()) {};

  const vec3d& direction () const {return direction_;};
  /** Returns the distribution function for a cone angular
      distribution, which is 1/omega, where omega is the total solid
      angle subtended by the cone, if we are within the cone and zero
      otherwise. */
  T_lambda distribution_function (const vec3d& dir, const vec3d&) const {
    T_lambda d(n); 
    d= cone_angular_distribution::distribution_function(dir, vec3d(), d);
    return d;
  };
  T_lambda distribution_function (const vec3d& dir, const vec3d&,
				  T_lambda& distr) const {
    // The first direction vector is the important one, since the
    // second one is only used when scattering
    const T_float costheta = dot(dir,direction_);
    distr = (costheta > cos_theta_) ? value_ : 0;
  };

  boost::shared_ptr<angular_distribution<blitz::Array<T,1> > > clone() const {
    return boost::shared_ptr<angular_distribution<blitz::Array<T,1> > >
      (new cone_angular_distribution<blitz::Array<T,1> >(*this));};
};


/** A composition of a ray, a normalization vector, an
    angular_distribution object, and a velocity vector. (Name is for
    historical reasons, the normalization vector was added later...)
    Is used as a return type when scattering or emitting rays, as all
    this information must be returned. */
template <typename T_lambda>
class mcrx::ray_distribution_pair {
private:
  boost::shared_ptr<ray<T_lambda> > rayptr_;
  T_lambda norm_;

  // can't keep a reference here because sometimes the distribution is
  // a temporary
  boost::shared_ptr<angular_distribution<T_lambda> > distr_;
  vec3d vel_;

public:
  ray_distribution_pair (const boost::shared_ptr<ray<T_lambda> >& f,
			 const T_lambda& n,
                         const angular_distribution<T_lambda>& s,
			 const vec3d& v=vec3d(0,0,0)):
    rayptr_ (f), norm_(make_thread_local_copy(n)), distr_ (s.clone()), vel_(v) {};

  boost::shared_ptr<ray<T_lambda> >& ray_pointer() {return rayptr_;};
  T_lambda& normalization() {return norm_;};
  const angular_distribution<T_lambda>& distribution() const {return *distr_;};
  const vec3d& velocity() const {return vel_;};
};
		      
  
#endif
